/*
A wierd thing: some modules cannot send data at some point, reporting failure in OnDataSent, even after cooling and rebooting
Pressing its shield with hand fix that
sticking a thick piece of paper on it also fix that
*/

#include "CommEspNow.hpp"

#include <esp_now.h>
//#include ""

#include <WiFi.h>
//#include <WiFi.h>
#include<Arduino.h>
#include <Preferences.h>
#include <esp_wifi.h>

const int peerNumMax = 8;
esp_now_peer_info_t peerInfo[peerNumMax];
uint32_t usingPeerMask = 0;
int rxPairIdx = 0;
char* RxBuffer;
int rxBufferSize;
uint8_t bufferFilled = 0;

uint8_t sequenceSegCountdown = 0;
clock_t sequenceRefreshTime = 0;
uint8_t sequencePairIdx = 0;
const int addressSize = 6;

#define NOW_TIME_OUT 100

Preferences preferencesComm;
Stream* espNowPrint;

//#define ESP_NOW_VERBOSE 

#define ESPNOW_NAMESPACE "ENCM" // esp now comm
#define CONFIG_PEER_MASK_ENTRY "PMSK" // peer mask
#define CONFIG_PEER_ENTRY_TEMPLATE "PR%d" // peer?

#define ESP_NOW_CH 1

void printMac(const uint8_t* mac_addr) {
	for (int idx = 0; idx < 5; idx++) {
		espNowPrint->print(mac_addr[idx], HEX);
		espNowPrint->print(",");
	}
	espNowPrint->print(mac_addr[5], HEX);
}

void OnDataSent(const uint8_t* mac_addr, esp_now_send_status_t status) {
	if (status == ESP_NOW_SEND_SUCCESS) {
#ifdef ESP_NOW_VERBOSE
		espNowPrint->println("Sent over NOW");
#endif
	}
	else
	{
		espNowPrint->print("Cannot sent over NOW to ");
		printMac(mac_addr);
		espNowPrint->print(": ");
		espNowPrint->println(status, HEX);
	}
}

void OnDataRecv(const uint8_t* mac, const uint8_t* incomingData, int len) {
#ifdef ESP_NOW_VERBOSE
	espNowPrint->println("Something over NOW");
#endif
	if (sequenceSegCountdown > 0) {
		bool allMatch = true;
		for (int addrIdx = 0; (addrIdx < 6) && allMatch; addrIdx++) {
			allMatch &= (mac[addrIdx] == peerInfo[sequencePairIdx].peer_addr[addrIdx]);
		}
		if (!allMatch) {
#ifdef ESP_NOW_VERBOSE
			espNowPrint->println("Interupt msg");
#endif
			return;
		}
	}
	for (int idx = 0; idx < peerNumMax; idx++) {
		if (bitRead(usingPeerMask, idx) == 0x01) {
			bool allMatch;
			for (int addrIdx = 0; (addrIdx < 6) && allMatch; addrIdx++) {
				allMatch &= (mac[addrIdx] == peerInfo[idx].peer_addr[addrIdx]);
			}
			if (allMatch) {
				rxPairIdx = idx;
				break;
			}
		}
	}
	if (incomingData[0] == 0x00) {
		if (incomingData[1] == 0x02) {
			sequencePairIdx = rxPairIdx;
			sequenceSegCountdown = incomingData[2];
			sequenceRefreshTime = clock();
			return;
		};
	}
	if (sequenceSegCountdown > 0) {
		sequenceSegCountdown--;
		sequenceRefreshTime = clock();
#ifdef ESP_NOW_VERBOSE
		espNowPrint->println(sequenceSegCountdown);
#endif
	}
	memcpy(RxBuffer + bufferFilled, incomingData, len);
	bufferFilled += len;
}

void savePeer(int index) {
	espNowPrint->println("Saving peers");
	if (preferencesComm.begin(ESPNOW_NAMESPACE, false)) {
		if (index < 0) {
			preferencesComm.putUInt(CONFIG_PEER_MASK_ENTRY, usingPeerMask);
			espNowPrint->println("Peer mask saved\r\n");
		}
		else {
			char peerEntryBuf[16];
			sprintf(peerEntryBuf, CONFIG_PEER_ENTRY_TEMPLATE, index);
			int putLen = preferencesComm.putBytes(peerEntryBuf, peerInfo[index].peer_addr, 6);
			if (putLen == 6) {
				espNowPrint->println("Peer " + String(index) + " saved");
				preferencesComm.putUInt(CONFIG_PEER_MASK_ENTRY, usingPeerMask);
				espNowPrint->println("Peer mask saved\r\n");
			}
			else {
				espNowPrint->print("Save peer error: ");
				espNowPrint->println(putLen);
			}
		}
	}
	else {
		espNowPrint->println("No config");
	}
	preferencesComm.end();
}

char* loadPeers(char* buffer) {
	espNowPrint->print("Loading peers ");
	if (preferencesComm.begin(ESPNOW_NAMESPACE, true)) {
		usingPeerMask = preferencesComm.getUInt(CONFIG_PEER_MASK_ENTRY, 0);

		espNowPrint->print("Peer mask: 0x");
		espNowPrint->println(usingPeerMask, HEX);

		char peerEntryBuf[16];
		byte peerMacBuffer[6];
		for (int idx = 0; idx < peerNumMax; idx++) {
			if (bitRead(usingPeerMask, idx) == 0x01) {
				sprintf(peerEntryBuf, CONFIG_PEER_ENTRY_TEMPLATE, idx);
				int readLen = preferencesComm.getBytes(peerEntryBuf, peerMacBuffer, 6);
				if (readLen == 6) {
					CommEspNow::addPeer(idx, peerMacBuffer, false);
					espNowPrint->println("Load peer " + String(idx) + ": ");
					espNowPrint->print("  ");
					printMac(peerMacBuffer);
					espNowPrint->println("");
				}
			}
		}
	}
	else {
		espNowPrint->println("No config");
	}
	preferencesComm.end();
	return buffer;
}

int loadStartupMode(char* bufferRef) {
	uint8_t ret = 255;
	espNowPrint->print("Loading startup mode ");
	if (preferencesComm.begin(ESPNOW_NAMESPACE, true)) {
		ret = preferencesComm.getUChar("StartUp", ret);
		espNowPrint->println(ret);
	}
	else {
		espNowPrint->println("No config");
	}
	preferencesComm.end();
	return ret;
}

int CommEspNow::init(void* buffer, int size, esp_now_recv_cb cb, Stream* stm) {
	//espNowPrint->begin(115200);
	if (stm == nullptr) {
		espNowPrint = &Serial;
	}
	else {
		espNowPrint = stm;
	}
	RxBuffer = (char*)buffer;
	rxBufferSize = size;
	esp_err_t err = ESP_OK;
	// Set device as a Wi-Fi Station
	WiFi.mode(WIFI_STA);
	//auto stat = WiFi.begin();
	//espNowPrint->print("Wifi stat: ");
	//espNowPrint->println(stat);

	// https://esp32.com/viewtopic.php?t=14542
	//wifi_config_t wifi_config = {
	//	.sta = {
	//	.channel = ESP_NOW_CH,
	//	},
	//};
	//ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
	//ESP_ERROR_CHECK(esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config));

	espNowPrint->println("Wifi mac");
	espNowPrint->println(WiFi.macAddress());

	// Init ESP-NOW
	if (esp_now_init() != ESP_OK) {
		espNowPrint->println("Error initializing ESP-NOW");
		return 0;
	}

	err = esp_wifi_start();
	if (err != ESP_OK) {
		espNowPrint->print("esp_wifi_start failed: ");
		espNowPrint->println(err, HEX);
	}

	// waiting for wifi start is actuall complicated
	// example here 
	// https://github.com/espressif/esp-idf/blob/46acfdce969f03c02b001fe4d24fa9e98f6adc5e/examples/wifi/getting_started/station/main/station_example_main.c

	//for (int idx = 0; idx < 10; idx++) {
	//	auto wifiStat = WiFi.status();
	//	espNowPrint->print("Wifi status: ");
	//	espNowPrint->println(wifiStat);

	//	err = esp_wifi_set_channel(ESP_NOW_CH, WIFI_SECOND_CHAN_NONE);
	//	if (err != ESP_OK) {
	//		espNowPrint->print("esp_wifi_set_channel failed: ");
	//		espNowPrint->println(err, HEX);
	//		espNowPrint->println("Retry");
	//	}
	//	else {
	//		if (WiFi.channel() != ESP_NOW_CH) {
	//			espNowPrint->println("Current channel not expected. Retry.");
	//		}
	//	}
	//	delay(1000);
	//}
	//esp_wifi_config_espnow_rate(WIFI_IF_STA, 100000);

	espNowPrint->println("Use default wifi channel");
	espNowPrint->println("Wifi channel");
	espNowPrint->println(WiFi.channel());

	if (esp_now_register_send_cb(OnDataSent) != ESP_OK) {
		espNowPrint->println("Error registering send cb");
	}
	if (cb != NULL) {
		if (esp_now_register_recv_cb(cb) != ESP_OK) {
			espNowPrint->println("Error registering receive cb");
		}
	}
	else {
		if (esp_now_register_recv_cb(OnDataRecv) != ESP_OK) {
			espNowPrint->println("Error registering receive cb");
		}
	}

	char* bufferRef = nullptr;
	bufferRef = loadPeers(bufferRef);
	return loadStartupMode(bufferRef);
}

void CommEspNow::setStartupMode(int val) {
	espNowPrint->print("Saving startup mode ");
	espNowPrint->println(val);
	if (preferencesComm.begin(ESPNOW_NAMESPACE, false)) {
		preferencesComm.putUChar("StartUp", val);
	}
	else {
		espNowPrint->println("No config");
	}
	preferencesComm.end();
}

int CommEspNow::clearPeers() {
	usingPeerMask = 0;
	savePeer(-1);
	return 1;
}

int CommEspNow::addPeer(int index, uint8_t* address, bool save) {
	if (index < 0) {
		for (int idx = 0; idx < peerNumMax; idx++) {
			if (bitRead(usingPeerMask, idx) != 0x01) {
				index = idx;
				break;
			}
		}
	}
	if (index < 0) {
		espNowPrint->print("Peer list full\r\n");
		espNowPrint->flush();
		return -1;
	}
	int pIdx = index;
	usingPeerMask |= bit(pIdx);
	memcpy(peerInfo[pIdx].peer_addr, address, addressSize);
	peerInfo[pIdx].ifidx = WIFI_IF_STA;
	peerInfo[pIdx].channel = ESP_NOW_CH;
	peerInfo[pIdx].encrypt = false;

	auto flag = esp_now_add_peer(peerInfo + pIdx);
	if (flag != ESP_OK) {
		espNowPrint->print("Failed to add peer # ");
		espNowPrint->print(pIdx);
		espNowPrint->print(": ");
		espNowPrint->println(flag);
		return -1;
	}

	if (save) {
		savePeer(index);
	}

	return pIdx;
}

void CommEspNow::sendStr(const char* msg) {
	sendStr(0, msg);
}

int CommEspNow::sendStr(int idx, const char* msg) {
	return sendStream(idx, (uint8_t*)msg, strlen(msg) + 1);
}

int CommEspNow::sendStream(int idx, const void* dat, int len) {
	if (bitRead(usingPeerMask, idx) == 0) {
		return ESP_ERR_INVALID_MAC;
	}
	esp_err_t result = esp_now_send(peerInfo[idx].peer_addr, (uint8_t*)dat, len);
	if (result == ESP_OK) {
#ifdef ESP_NOW_VERBOSE
		espNowPrint->println("Sent with success");
#endif
	}
	else
	{
		espNowPrint->print("Error sending the data: ");
		espNowPrint->println(result, HEX);
	}
	return result;
}

int CommEspNow::receiveStr() {
	if (bufferFilled == 0) {
		return 0;
	}
	if (sequenceSegCountdown > 0) {
		if (clock() - sequenceRefreshTime < NOW_TIME_OUT) {
			return 0;
		}
		else {
			sequenceSegCountdown = 0;
		}
	}
	int ret = bufferFilled;
	//memcpy(msg, RxBuffer, bufferFilled);
	bufferFilled = 0;
	return ret;
}

void CommEspNow::printStatus() {
	espNowPrint->println("Wifi mac");
	espNowPrint->println(WiFi.macAddress());
	espNowPrint->print("Peer mask: 0x");
	espNowPrint->println(usingPeerMask, HEX);
	for (int idx = 0; idx < peerNumMax; idx++) {
		if (bitRead(usingPeerMask, idx) == 0x01) {
			espNowPrint->print("  Peer #");
			espNowPrint->print(idx);
			espNowPrint->print(" mac ");
			printMac(peerInfo[idx].peer_addr);
			espNowPrint->println("");
		}
	}
	espNowPrint->println("End");
}

// not tested.
int CommEspNow::sendLineByLine(int idx, const char* msgLines, int timeout) {
	int count = 0;
	for (const char* c = msgLines; c[0] != '\0'; c++) {
		if (c[0] == '\n') {
			count++;
		}
	}

	uint8_t flags[3];
	flags[0] = 0;
	flags[1] = 0x02;
	flags[2] = count;
	sendStream(idx, flags, 3);
	const char* lineStart = msgLines;
	for (int lineIdx = 0; lineIdx < count; lineIdx++) {
		const char* lineEnd = lineStart;
		while (lineEnd[0] != '\n') {
			lineEnd++;
		}

		int exBit = (lineIdx == count - 1) ? 1 : 0;
		sendStream(idx, (uint8_t*)lineStart, lineEnd - lineStart + 1 + exBit);
		lineStart = lineEnd + 1;
	}
	return 0;
}
